from langdev_helper.llm.lcex import llm_lcex as model

from tool import get_context, cite_context_sources

import operator
from typing import Annotated, Sequence
from typing_extensions import TypedDict
from langchain_core.messages import BaseMessage

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]


# from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolNode

# model = ChatOpenAI(model="gpt-4o", temperature=0)

# 定义决定是否继续的函数
def should_continue(state, config):
    messages = state["messages"]
    last_message = messages[-1]
    # 如果没有函数调用，则结束
    if not last_message.tool_calls:
        return "end"
    # 否则继续
    else:
        return "continue"

tools = [get_context, cite_context_sources]

# 定义调用模型的函数
def call_model(state, config):
    messages = state["messages"]
    model_with_tools = model.bind_tools(tools)
    response = model_with_tools.invoke(messages)
    # 返回一个列表，因为这将添加到现有列表中
    return {"messages": [response]}

# ToolNode 将自动处理将状态注入到工具中
tool_node = ToolNode(tools)



from langgraph.graph import END, START, StateGraph

# 定义一个新的图
workflow = StateGraph(AgentState)

# 定义两个节点以循环使用
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

# 将入口点设置为 `agent`
# 这意味着该节点是第一个被调用的
workflow.add_edge(START, "agent")

# 现在我们添加一个条件边
workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "continue": "action",
        "end": END,
    },
)

# 现在我们添加一条从 `tools` 到 `agent` 的普通边
# 这意味着在调用 `tools` 后，下一个调用将是 `agent` 节点
workflow.add_edge("action", "agent")

# 最后，编译它！
# 这将其编译为 LangChain 可运行对象，
# 意味着用户可以像使用任何其他可运行对象一样使用它
graph = workflow.compile()

print(graph.get_graph().draw_mermaid())



# ==== use
from langchain_core.messages import HumanMessage

messages = [HumanMessage("关于 FooBar 的最新消息是什么？")]
for output in graph.stream({"messages": messages}):
    # stream() 生成以节点名称为键的输出字典
    for key, value in output.items():
        print(f"来自节点 '{key}' 的输出:")
        print("---")
        print(value)
        messages.extend(value["messages"])
    print("\n---\n")


# ==== use again
messages.append(HumanMessage("你从哪里获得这些信息？"))
for output in graph.stream({"messages": messages}):
    # stream() 生成以节点名称为键的输出字典
    for key, value in output.items():
        print(f"来自节点 '{key}' 的输出:")
        print("---")
        print(value)
    print("\n---\n")
